
<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Transitional//EN"
  "http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd">

<html xmlns="http://www.w3.org/1999/xhtml" lang="en">
  <head>
    <meta http-equiv="X-UA-Compatible" content="IE=Edge" />
    <meta http-equiv="Content-Type" content="text/html; charset=utf-8" />
    <title>tensorflowonspark.pipeline &#8212; TensorFlowOnSpark 1.3.3 documentation</title>
    <link rel="stylesheet" href="../../_static/classic.css" type="text/css" />
    <link rel="stylesheet" href="../../_static/pygments.css" type="text/css" />
    <script type="text/javascript" id="documentation_options" data-url_root="../../" src="../../_static/documentation_options.js"></script>
    <script type="text/javascript" src="../../_static/jquery.js"></script>
    <script type="text/javascript" src="../../_static/underscore.js"></script>
    <script type="text/javascript" src="../../_static/doctools.js"></script>
    <link rel="index" title="Index" href="../../genindex.html" />
    <link rel="search" title="Search" href="../../search.html" /> 
  </head><body>
    <div class="related" role="navigation" aria-label="related navigation">
      <h3>Navigation</h3>
      <ul>
        <li class="right" style="margin-right: 10px">
          <a href="../../genindex.html" title="General Index"
             accesskey="I">index</a></li>
        <li class="right" >
          <a href="../../py-modindex.html" title="Python Module Index"
             >modules</a> |</li>
        <li class="nav-item nav-item-0"><a href="../../index.html">TensorFlowOnSpark 1.3.3 documentation</a> &#187;</li>
          <li class="nav-item nav-item-1"><a href="../index.html" accesskey="U">Module code</a> &#187;</li> 
      </ul>
    </div>  

    <div class="document">
      <div class="documentwrapper">
        <div class="bodywrapper">
          <div class="body" role="main">
            
  <h1>Source code for tensorflowonspark.pipeline</h1><div class="highlight"><pre>
<span></span><span class="c1"># Copyright 2017 Yahoo Inc.</span>
<span class="c1"># Licensed under the terms of the Apache 2.0 license.</span>
<span class="c1"># Please see LICENSE file in the project root for terms.</span>
<span class="sd">&quot;&quot;&quot;This module extends the TensorFlowOnSpark API to support Spark ML Pipelines.</span>

<span class="sd">It provides a TFEstimator class to fit a TFModel using TensorFlow.  The TFEstimator will actually spawn a TensorFlowOnSpark cluster</span>
<span class="sd">to conduct distributed training, but due to architectural limitations, the TFModel will only run single-node TensorFlow instances</span>
<span class="sd">when inferencing on the executors.  The executors will run in parallel, but the TensorFlow model must fit in the memory</span>
<span class="sd">of each executor.</span>

<span class="sd">There is also an option to provide a separate &quot;export&quot; function, which allows users to export a different graph for inferencing vs. training.</span>
<span class="sd">This is useful when the training graph uses InputMode.TENSORFLOW with queue_runners, but the inferencing graph needs placeholders.</span>
<span class="sd">And this is especially useful for exporting saved_models for TensorFlow Serving.</span>
<span class="sd">&quot;&quot;&quot;</span>

<span class="kn">from</span> <span class="nn">__future__</span> <span class="k">import</span> <span class="n">absolute_import</span>
<span class="kn">from</span> <span class="nn">__future__</span> <span class="k">import</span> <span class="n">division</span>
<span class="kn">from</span> <span class="nn">__future__</span> <span class="k">import</span> <span class="n">print_function</span>

<span class="kn">from</span> <span class="nn">pyspark.context</span> <span class="k">import</span> <span class="n">SparkContext</span>
<span class="kn">from</span> <span class="nn">pyspark.ml.param.shared</span> <span class="k">import</span> <span class="n">Param</span><span class="p">,</span> <span class="n">Params</span><span class="p">,</span> <span class="n">TypeConverters</span>
<span class="kn">from</span> <span class="nn">pyspark.ml.pipeline</span> <span class="k">import</span> <span class="n">Estimator</span><span class="p">,</span> <span class="n">Model</span>
<span class="kn">from</span> <span class="nn">pyspark.sql</span> <span class="k">import</span> <span class="n">Row</span><span class="p">,</span> <span class="n">SparkSession</span>

<span class="kn">import</span> <span class="nn">tensorflow</span> <span class="k">as</span> <span class="nn">tf</span>
<span class="kn">from</span> <span class="nn">tensorflow.contrib.saved_model.python.saved_model</span> <span class="k">import</span> <span class="n">reader</span><span class="p">,</span> <span class="n">signature_def_utils</span>
<span class="kn">from</span> <span class="nn">tensorflow.python.saved_model</span> <span class="k">import</span> <span class="n">loader</span>
<span class="kn">from</span> <span class="nn">.</span> <span class="k">import</span> <span class="n">TFCluster</span><span class="p">,</span> <span class="n">gpu_info</span><span class="p">,</span> <span class="n">dfutil</span>

<span class="kn">import</span> <span class="nn">argparse</span>
<span class="kn">import</span> <span class="nn">copy</span>
<span class="kn">import</span> <span class="nn">logging</span>
<span class="kn">import</span> <span class="nn">os</span>
<span class="kn">import</span> <span class="nn">subprocess</span>
<span class="kn">import</span> <span class="nn">sys</span>


<span class="c1"># TensorFlowOnSpark Params</span>

<div class="viewcode-block" id="TFTypeConverters"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.TFTypeConverters">[docs]</a><span class="k">class</span> <span class="nc">TFTypeConverters</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;Custom DataFrame TypeConverter for dictionary types (since this is not provided by Spark core).&quot;&quot;&quot;</span>
<div class="viewcode-block" id="TFTypeConverters.toDict"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.TFTypeConverters.toDict">[docs]</a>  <span class="nd">@staticmethod</span>
  <span class="k">def</span> <span class="nf">toDict</span><span class="p">(</span><span class="n">value</span><span class="p">):</span>
    <span class="k">if</span> <span class="nb">type</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> <span class="o">==</span> <span class="nb">dict</span><span class="p">:</span>
      <span class="k">return</span> <span class="n">value</span>
    <span class="k">else</span><span class="p">:</span>
      <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;Could not convert </span><span class="si">%s</span><span class="s2"> to OrderedDict&quot;</span> <span class="o">%</span> <span class="n">value</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="HasBatchSize"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasBatchSize">[docs]</a><span class="k">class</span> <span class="nc">HasBatchSize</span><span class="p">(</span><span class="n">Params</span><span class="p">):</span>
  <span class="n">batch_size</span> <span class="o">=</span> <span class="n">Param</span><span class="p">(</span><span class="n">Params</span><span class="o">.</span><span class="n">_dummy</span><span class="p">(),</span> <span class="s2">&quot;batch_size&quot;</span><span class="p">,</span> <span class="s2">&quot;Number of records per batch&quot;</span><span class="p">,</span> <span class="n">typeConverter</span><span class="o">=</span><span class="n">TypeConverters</span><span class="o">.</span><span class="n">toInt</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">HasBatchSize</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

<div class="viewcode-block" id="HasBatchSize.setBatchSize"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasBatchSize.setBatchSize">[docs]</a>  <span class="k">def</span> <span class="nf">setBatchSize</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_set</span><span class="p">(</span><span class="n">batch_size</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasBatchSize.getBatchSize"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasBatchSize.getBatchSize">[docs]</a>  <span class="k">def</span> <span class="nf">getBatchSize</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">batch_size</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="HasClusterSize"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasClusterSize">[docs]</a><span class="k">class</span> <span class="nc">HasClusterSize</span><span class="p">(</span><span class="n">Params</span><span class="p">):</span>
  <span class="n">cluster_size</span> <span class="o">=</span> <span class="n">Param</span><span class="p">(</span><span class="n">Params</span><span class="o">.</span><span class="n">_dummy</span><span class="p">(),</span> <span class="s2">&quot;cluster_size&quot;</span><span class="p">,</span> <span class="s2">&quot;Number of nodes in the cluster&quot;</span><span class="p">,</span> <span class="n">typeConverter</span><span class="o">=</span><span class="n">TypeConverters</span><span class="o">.</span><span class="n">toInt</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">HasClusterSize</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

<div class="viewcode-block" id="HasClusterSize.setClusterSize"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasClusterSize.setClusterSize">[docs]</a>  <span class="k">def</span> <span class="nf">setClusterSize</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_set</span><span class="p">(</span><span class="n">cluster_size</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasClusterSize.getClusterSize"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasClusterSize.getClusterSize">[docs]</a>  <span class="k">def</span> <span class="nf">getClusterSize</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cluster_size</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="HasEpochs"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasEpochs">[docs]</a><span class="k">class</span> <span class="nc">HasEpochs</span><span class="p">(</span><span class="n">Params</span><span class="p">):</span>
  <span class="n">epochs</span> <span class="o">=</span> <span class="n">Param</span><span class="p">(</span><span class="n">Params</span><span class="o">.</span><span class="n">_dummy</span><span class="p">(),</span> <span class="s2">&quot;epochs&quot;</span><span class="p">,</span> <span class="s2">&quot;Number of epochs to train&quot;</span><span class="p">,</span> <span class="n">typeConverter</span><span class="o">=</span><span class="n">TypeConverters</span><span class="o">.</span><span class="n">toInt</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">HasEpochs</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

<div class="viewcode-block" id="HasEpochs.setEpochs"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasEpochs.setEpochs">[docs]</a>  <span class="k">def</span> <span class="nf">setEpochs</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_set</span><span class="p">(</span><span class="n">epochs</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasEpochs.getEpochs"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasEpochs.getEpochs">[docs]</a>  <span class="k">def</span> <span class="nf">getEpochs</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">epochs</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="HasInputMapping"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasInputMapping">[docs]</a><span class="k">class</span> <span class="nc">HasInputMapping</span><span class="p">(</span><span class="n">Params</span><span class="p">):</span>
  <span class="n">input_mapping</span> <span class="o">=</span> <span class="n">Param</span><span class="p">(</span><span class="n">Params</span><span class="o">.</span><span class="n">_dummy</span><span class="p">(),</span> <span class="s2">&quot;input_mapping&quot;</span><span class="p">,</span> <span class="s2">&quot;Mapping of input DataFrame column to input tensor&quot;</span><span class="p">,</span> <span class="n">typeConverter</span><span class="o">=</span><span class="n">TFTypeConverters</span><span class="o">.</span><span class="n">toDict</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">HasInputMapping</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

<div class="viewcode-block" id="HasInputMapping.setInputMapping"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasInputMapping.setInputMapping">[docs]</a>  <span class="k">def</span> <span class="nf">setInputMapping</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_set</span><span class="p">(</span><span class="n">input_mapping</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasInputMapping.getInputMapping"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasInputMapping.getInputMapping">[docs]</a>  <span class="k">def</span> <span class="nf">getInputMapping</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">input_mapping</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="HasInputMode"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasInputMode">[docs]</a><span class="k">class</span> <span class="nc">HasInputMode</span><span class="p">(</span><span class="n">Params</span><span class="p">):</span>
  <span class="n">input_mode</span> <span class="o">=</span> <span class="n">Param</span><span class="p">(</span><span class="n">Params</span><span class="o">.</span><span class="n">_dummy</span><span class="p">(),</span> <span class="s2">&quot;input_mode&quot;</span><span class="p">,</span> <span class="s2">&quot;Input data feeding mode (0=TENSORFLOW, 1=SPARK)&quot;</span><span class="p">,</span> <span class="n">typeConverter</span><span class="o">=</span><span class="n">TypeConverters</span><span class="o">.</span><span class="n">toInt</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">HasInputMode</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

<div class="viewcode-block" id="HasInputMode.setInputMode"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasInputMode.setInputMode">[docs]</a>  <span class="k">def</span> <span class="nf">setInputMode</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_set</span><span class="p">(</span><span class="n">input_mode</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasInputMode.getInputMode"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasInputMode.getInputMode">[docs]</a>  <span class="k">def</span> <span class="nf">getInputMode</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">input_mode</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="HasModelDir"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasModelDir">[docs]</a><span class="k">class</span> <span class="nc">HasModelDir</span><span class="p">(</span><span class="n">Params</span><span class="p">):</span>
  <span class="n">model_dir</span> <span class="o">=</span> <span class="n">Param</span><span class="p">(</span><span class="n">Params</span><span class="o">.</span><span class="n">_dummy</span><span class="p">(),</span> <span class="s2">&quot;model_dir&quot;</span><span class="p">,</span> <span class="s2">&quot;Path to save/load model checkpoints&quot;</span><span class="p">,</span> <span class="n">typeConverter</span><span class="o">=</span><span class="n">TypeConverters</span><span class="o">.</span><span class="n">toString</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">HasModelDir</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

<div class="viewcode-block" id="HasModelDir.setModelDir"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasModelDir.setModelDir">[docs]</a>  <span class="k">def</span> <span class="nf">setModelDir</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_set</span><span class="p">(</span><span class="n">model_dir</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasModelDir.getModelDir"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasModelDir.getModelDir">[docs]</a>  <span class="k">def</span> <span class="nf">getModelDir</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">model_dir</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="HasNumPS"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasNumPS">[docs]</a><span class="k">class</span> <span class="nc">HasNumPS</span><span class="p">(</span><span class="n">Params</span><span class="p">):</span>
  <span class="n">num_ps</span> <span class="o">=</span> <span class="n">Param</span><span class="p">(</span><span class="n">Params</span><span class="o">.</span><span class="n">_dummy</span><span class="p">(),</span> <span class="s2">&quot;num_ps&quot;</span><span class="p">,</span> <span class="s2">&quot;Number of PS nodes in cluster&quot;</span><span class="p">,</span> <span class="n">typeConverter</span><span class="o">=</span><span class="n">TypeConverters</span><span class="o">.</span><span class="n">toInt</span><span class="p">)</span>
  <span class="n">driver_ps_nodes</span> <span class="o">=</span> <span class="n">Param</span><span class="p">(</span><span class="n">Params</span><span class="o">.</span><span class="n">_dummy</span><span class="p">(),</span> <span class="s2">&quot;driver_ps_nodes&quot;</span><span class="p">,</span> <span class="s2">&quot;Run PS nodes on driver locally&quot;</span><span class="p">,</span> <span class="n">typeConverter</span><span class="o">=</span><span class="n">TypeConverters</span><span class="o">.</span><span class="n">toBoolean</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">HasNumPS</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

<div class="viewcode-block" id="HasNumPS.setNumPS"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasNumPS.setNumPS">[docs]</a>  <span class="k">def</span> <span class="nf">setNumPS</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_set</span><span class="p">(</span><span class="n">num_ps</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasNumPS.getNumPS"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasNumPS.getNumPS">[docs]</a>  <span class="k">def</span> <span class="nf">getNumPS</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">num_ps</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasNumPS.setDriverPSNodes"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasNumPS.setDriverPSNodes">[docs]</a>  <span class="k">def</span> <span class="nf">setDriverPSNodes</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_set</span><span class="p">(</span><span class="n">driver_ps_nodes</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasNumPS.getDriverPSNodes"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasNumPS.getDriverPSNodes">[docs]</a>  <span class="k">def</span> <span class="nf">getDriverPSNodes</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">driver_ps_nodes</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="HasOutputMapping"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasOutputMapping">[docs]</a><span class="k">class</span> <span class="nc">HasOutputMapping</span><span class="p">(</span><span class="n">Params</span><span class="p">):</span>
  <span class="n">output_mapping</span> <span class="o">=</span> <span class="n">Param</span><span class="p">(</span><span class="n">Params</span><span class="o">.</span><span class="n">_dummy</span><span class="p">(),</span> <span class="s2">&quot;output_mapping&quot;</span><span class="p">,</span> <span class="s2">&quot;Mapping of output tensor to output DataFrame column&quot;</span><span class="p">,</span> <span class="n">typeConverter</span><span class="o">=</span><span class="n">TFTypeConverters</span><span class="o">.</span><span class="n">toDict</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">HasOutputMapping</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

<div class="viewcode-block" id="HasOutputMapping.setOutputMapping"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasOutputMapping.setOutputMapping">[docs]</a>  <span class="k">def</span> <span class="nf">setOutputMapping</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_set</span><span class="p">(</span><span class="n">output_mapping</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasOutputMapping.getOutputMapping"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasOutputMapping.getOutputMapping">[docs]</a>  <span class="k">def</span> <span class="nf">getOutputMapping</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">output_mapping</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="HasProtocol"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasProtocol">[docs]</a><span class="k">class</span> <span class="nc">HasProtocol</span><span class="p">(</span><span class="n">Params</span><span class="p">):</span>
  <span class="n">protocol</span> <span class="o">=</span> <span class="n">Param</span><span class="p">(</span><span class="n">Params</span><span class="o">.</span><span class="n">_dummy</span><span class="p">(),</span> <span class="s2">&quot;protocol&quot;</span><span class="p">,</span> <span class="s2">&quot;Network protocol for Tensorflow (grpc|rdma)&quot;</span><span class="p">,</span> <span class="n">typeConverter</span><span class="o">=</span><span class="n">TypeConverters</span><span class="o">.</span><span class="n">toString</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">HasProtocol</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

<div class="viewcode-block" id="HasProtocol.setProtocol"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasProtocol.setProtocol">[docs]</a>  <span class="k">def</span> <span class="nf">setProtocol</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_set</span><span class="p">(</span><span class="n">protocol</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasProtocol.getProtocol"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasProtocol.getProtocol">[docs]</a>  <span class="k">def</span> <span class="nf">getProtocol</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">protocol</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="HasReaders"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasReaders">[docs]</a><span class="k">class</span> <span class="nc">HasReaders</span><span class="p">(</span><span class="n">Params</span><span class="p">):</span>
  <span class="n">readers</span> <span class="o">=</span> <span class="n">Param</span><span class="p">(</span><span class="n">Params</span><span class="o">.</span><span class="n">_dummy</span><span class="p">(),</span> <span class="s2">&quot;readers&quot;</span><span class="p">,</span> <span class="s2">&quot;number of reader/enqueue threads&quot;</span><span class="p">,</span> <span class="n">typeConverter</span><span class="o">=</span><span class="n">TypeConverters</span><span class="o">.</span><span class="n">toInt</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">HasReaders</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

<div class="viewcode-block" id="HasReaders.setReaders"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasReaders.setReaders">[docs]</a>  <span class="k">def</span> <span class="nf">setReaders</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_set</span><span class="p">(</span><span class="n">readers</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasReaders.getReaders"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasReaders.getReaders">[docs]</a>  <span class="k">def</span> <span class="nf">getReaders</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">readers</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="HasSteps"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasSteps">[docs]</a><span class="k">class</span> <span class="nc">HasSteps</span><span class="p">(</span><span class="n">Params</span><span class="p">):</span>
  <span class="n">steps</span> <span class="o">=</span> <span class="n">Param</span><span class="p">(</span><span class="n">Params</span><span class="o">.</span><span class="n">_dummy</span><span class="p">(),</span> <span class="s2">&quot;steps&quot;</span><span class="p">,</span> <span class="s2">&quot;Maximum number of steps to train&quot;</span><span class="p">,</span> <span class="n">typeConverter</span><span class="o">=</span><span class="n">TypeConverters</span><span class="o">.</span><span class="n">toInt</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">HasSteps</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

<div class="viewcode-block" id="HasSteps.setSteps"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasSteps.setSteps">[docs]</a>  <span class="k">def</span> <span class="nf">setSteps</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_set</span><span class="p">(</span><span class="n">steps</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasSteps.getSteps"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasSteps.getSteps">[docs]</a>  <span class="k">def</span> <span class="nf">getSteps</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">steps</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="HasTensorboard"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasTensorboard">[docs]</a><span class="k">class</span> <span class="nc">HasTensorboard</span><span class="p">(</span><span class="n">Params</span><span class="p">):</span>
  <span class="n">tensorboard</span> <span class="o">=</span> <span class="n">Param</span><span class="p">(</span><span class="n">Params</span><span class="o">.</span><span class="n">_dummy</span><span class="p">(),</span> <span class="s2">&quot;tensorboard&quot;</span><span class="p">,</span> <span class="s2">&quot;Launch tensorboard process&quot;</span><span class="p">,</span> <span class="n">typeConverter</span><span class="o">=</span><span class="n">TypeConverters</span><span class="o">.</span><span class="n">toBoolean</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">HasTensorboard</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

<div class="viewcode-block" id="HasTensorboard.setTensorboard"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasTensorboard.setTensorboard">[docs]</a>  <span class="k">def</span> <span class="nf">setTensorboard</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_set</span><span class="p">(</span><span class="n">tensorboard</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasTensorboard.getTensorboard"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasTensorboard.getTensorboard">[docs]</a>  <span class="k">def</span> <span class="nf">getTensorboard</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">tensorboard</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="HasTFRecordDir"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasTFRecordDir">[docs]</a><span class="k">class</span> <span class="nc">HasTFRecordDir</span><span class="p">(</span><span class="n">Params</span><span class="p">):</span>
  <span class="n">tfrecord_dir</span> <span class="o">=</span> <span class="n">Param</span><span class="p">(</span><span class="n">Params</span><span class="o">.</span><span class="n">_dummy</span><span class="p">(),</span> <span class="s2">&quot;tfrecord_dir&quot;</span><span class="p">,</span> <span class="s2">&quot;Path to temporarily export a DataFrame as TFRecords (for InputMode.TENSORFLOW apps)&quot;</span><span class="p">,</span> <span class="n">typeConverter</span><span class="o">=</span><span class="n">TypeConverters</span><span class="o">.</span><span class="n">toString</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">HasTFRecordDir</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

<div class="viewcode-block" id="HasTFRecordDir.setTFRecordDir"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasTFRecordDir.setTFRecordDir">[docs]</a>  <span class="k">def</span> <span class="nf">setTFRecordDir</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_set</span><span class="p">(</span><span class="n">tfrecord_dir</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasTFRecordDir.getTFRecordDir"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasTFRecordDir.getTFRecordDir">[docs]</a>  <span class="k">def</span> <span class="nf">getTFRecordDir</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">tfrecord_dir</span><span class="p">)</span></div></div>


<span class="c1"># SavedModelBuilder Params</span>

<div class="viewcode-block" id="HasExportDir"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasExportDir">[docs]</a><span class="k">class</span> <span class="nc">HasExportDir</span><span class="p">(</span><span class="n">Params</span><span class="p">):</span>
  <span class="n">export_dir</span> <span class="o">=</span> <span class="n">Param</span><span class="p">(</span><span class="n">Params</span><span class="o">.</span><span class="n">_dummy</span><span class="p">(),</span> <span class="s2">&quot;export_dir&quot;</span><span class="p">,</span> <span class="s2">&quot;Directory to export saved_model&quot;</span><span class="p">,</span> <span class="n">typeConverter</span><span class="o">=</span><span class="n">TypeConverters</span><span class="o">.</span><span class="n">toString</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">HasExportDir</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

<div class="viewcode-block" id="HasExportDir.setExportDir"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasExportDir.setExportDir">[docs]</a>  <span class="k">def</span> <span class="nf">setExportDir</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_set</span><span class="p">(</span><span class="n">export_dir</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasExportDir.getExportDir"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasExportDir.getExportDir">[docs]</a>  <span class="k">def</span> <span class="nf">getExportDir</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">export_dir</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="HasSignatureDefKey"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasSignatureDefKey">[docs]</a><span class="k">class</span> <span class="nc">HasSignatureDefKey</span><span class="p">(</span><span class="n">Params</span><span class="p">):</span>
  <span class="n">signature_def_key</span> <span class="o">=</span> <span class="n">Param</span><span class="p">(</span><span class="n">Params</span><span class="o">.</span><span class="n">_dummy</span><span class="p">(),</span> <span class="s2">&quot;signature_def_key&quot;</span><span class="p">,</span> <span class="s2">&quot;Identifier for a specific saved_model signature&quot;</span><span class="p">,</span> <span class="n">typeConverter</span><span class="o">=</span><span class="n">TypeConverters</span><span class="o">.</span><span class="n">toString</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">HasSignatureDefKey</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">_setDefault</span><span class="p">(</span><span class="n">signature_def_key</span><span class="o">=</span><span class="kc">None</span><span class="p">)</span>

<div class="viewcode-block" id="HasSignatureDefKey.setSignatureDefKey"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasSignatureDefKey.setSignatureDefKey">[docs]</a>  <span class="k">def</span> <span class="nf">setSignatureDefKey</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_set</span><span class="p">(</span><span class="n">signature_def_key</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasSignatureDefKey.getSignatureDefKey"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasSignatureDefKey.getSignatureDefKey">[docs]</a>  <span class="k">def</span> <span class="nf">getSignatureDefKey</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">signature_def_key</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="HasTagSet"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasTagSet">[docs]</a><span class="k">class</span> <span class="nc">HasTagSet</span><span class="p">(</span><span class="n">Params</span><span class="p">):</span>
  <span class="n">tag_set</span> <span class="o">=</span> <span class="n">Param</span><span class="p">(</span><span class="n">Params</span><span class="o">.</span><span class="n">_dummy</span><span class="p">(),</span> <span class="s2">&quot;tag_set&quot;</span><span class="p">,</span> <span class="s2">&quot;Comma-delimited list of tags identifying a saved_model metagraph&quot;</span><span class="p">,</span> <span class="n">typeConverter</span><span class="o">=</span><span class="n">TypeConverters</span><span class="o">.</span><span class="n">toString</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">HasTagSet</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

<div class="viewcode-block" id="HasTagSet.setTagSet"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasTagSet.setTagSet">[docs]</a>  <span class="k">def</span> <span class="nf">setTagSet</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_set</span><span class="p">(</span><span class="n">tag_set</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>

<div class="viewcode-block" id="HasTagSet.getTagSet"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.HasTagSet.getTagSet">[docs]</a>  <span class="k">def</span> <span class="nf">getTagSet</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">tag_set</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="Namespace"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.Namespace">[docs]</a><span class="k">class</span> <span class="nc">Namespace</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">  Utility class to convert dictionaries to Namespace-like objects.</span>

<span class="sd">  Based on https://docs.python.org/dev/library/types.html#types.SimpleNamespace</span>
<span class="sd">  &quot;&quot;&quot;</span>
  <span class="n">argv</span> <span class="o">=</span> <span class="kc">None</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">d</span><span class="p">):</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">d</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
      <span class="bp">self</span><span class="o">.</span><span class="n">argv</span> <span class="o">=</span> <span class="n">d</span>
    <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">d</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
      <span class="bp">self</span><span class="o">.</span><span class="vm">__dict__</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">d</span><span class="p">)</span>
    <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">d</span><span class="p">,</span> <span class="n">argparse</span><span class="o">.</span><span class="n">Namespace</span><span class="p">):</span>
      <span class="bp">self</span><span class="o">.</span><span class="vm">__dict__</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="nb">vars</span><span class="p">(</span><span class="n">d</span><span class="p">))</span>
    <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">d</span><span class="p">,</span> <span class="n">Namespace</span><span class="p">):</span>
      <span class="bp">self</span><span class="o">.</span><span class="vm">__dict__</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">d</span><span class="o">.</span><span class="vm">__dict__</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
      <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;Unsupported Namespace args: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">d</span><span class="p">))</span>

  <span class="k">def</span> <span class="nf">__iter__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">argv</span><span class="p">:</span>
      <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">argv</span><span class="p">:</span>
        <span class="k">yield</span> <span class="n">item</span>
    <span class="k">else</span><span class="p">:</span>
      <span class="k">for</span> <span class="n">key</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__dict__</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
        <span class="k">yield</span> <span class="n">key</span>

  <span class="k">def</span> <span class="nf">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">argv</span><span class="p">:</span>
      <span class="k">return</span> <span class="s2">&quot;</span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">argv</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
      <span class="n">keys</span> <span class="o">=</span> <span class="nb">sorted</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="vm">__dict__</span><span class="p">)</span>
      <span class="n">items</span> <span class="o">=</span> <span class="p">(</span><span class="s2">&quot;</span><span class="si">{}</span><span class="s2">=</span><span class="si">{!r}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">k</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__dict__</span><span class="p">[</span><span class="n">k</span><span class="p">])</span> <span class="k">for</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">keys</span><span class="p">)</span>
      <span class="k">return</span> <span class="s2">&quot;</span><span class="si">{}</span><span class="s2">(</span><span class="si">{}</span><span class="s2">)&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span><span class="p">,</span> <span class="s2">&quot;, &quot;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">items</span><span class="p">))</span>

  <span class="k">def</span> <span class="nf">__eq__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">other</span><span class="p">):</span>
    <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">argv</span><span class="p">:</span>
      <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">argv</span> <span class="o">==</span> <span class="n">other</span>
    <span class="k">else</span><span class="p">:</span>
      <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__dict__</span> <span class="o">==</span> <span class="n">other</span><span class="o">.</span><span class="vm">__dict__</span></div>


<div class="viewcode-block" id="TFParams"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.TFParams">[docs]</a><span class="k">class</span> <span class="nc">TFParams</span><span class="p">(</span><span class="n">Params</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;Mix-in class to store namespace-style args and merge w/ SparkML-style params.&quot;&quot;&quot;</span>
  <span class="n">args</span> <span class="o">=</span> <span class="kc">None</span>

<div class="viewcode-block" id="TFParams.merge_args_params"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.TFParams.merge_args_params">[docs]</a>  <span class="k">def</span> <span class="nf">merge_args_params</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="n">local_args</span> <span class="o">=</span> <span class="n">copy</span><span class="o">.</span><span class="n">copy</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">args</span><span class="p">)</span>                 <span class="c1"># make a local copy of args</span>
    <span class="n">args_dict</span> <span class="o">=</span> <span class="nb">vars</span><span class="p">(</span><span class="n">local_args</span><span class="p">)</span>                      <span class="c1"># get dictionary view</span>
    <span class="k">for</span> <span class="n">p</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">params</span><span class="p">:</span>
      <span class="n">args_dict</span><span class="p">[</span><span class="n">p</span><span class="o">.</span><span class="n">name</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">getOrDefault</span><span class="p">(</span><span class="n">p</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>   <span class="c1"># update with params</span>
    <span class="k">return</span> <span class="n">local_args</span></div></div>


<div class="viewcode-block" id="TFEstimator"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.TFEstimator">[docs]</a><span class="k">class</span> <span class="nc">TFEstimator</span><span class="p">(</span><span class="n">Estimator</span><span class="p">,</span> <span class="n">TFParams</span><span class="p">,</span> <span class="n">HasInputMapping</span><span class="p">,</span>
                  <span class="n">HasClusterSize</span><span class="p">,</span> <span class="n">HasNumPS</span><span class="p">,</span> <span class="n">HasInputMode</span><span class="p">,</span> <span class="n">HasProtocol</span><span class="p">,</span> <span class="n">HasTensorboard</span><span class="p">,</span> <span class="n">HasModelDir</span><span class="p">,</span> <span class="n">HasExportDir</span><span class="p">,</span> <span class="n">HasTFRecordDir</span><span class="p">,</span>
                  <span class="n">HasBatchSize</span><span class="p">,</span> <span class="n">HasEpochs</span><span class="p">,</span> <span class="n">HasReaders</span><span class="p">,</span> <span class="n">HasSteps</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;Spark ML Estimator which launches a TensorFlowOnSpark cluster for distributed training.</span>

<span class="sd">  The columns of the DataFrame passed to the ``fit()`` method will be mapped to TensorFlow tensors according to the ``setInputMapping()`` method.</span>

<span class="sd">  If an ``export_fn`` was provided to the constructor, it will be run on a single executor immediately after the distributed training has completed.</span>
<span class="sd">  This allows users to export a TensorFlow saved_model with a different execution graph for inferencing, e.g. replacing an input graph of</span>
<span class="sd">  TFReaders and QueueRunners with Placeholders.</span>

<span class="sd">  For InputMode.TENSORFLOW, the input DataFrame will be exported as TFRecords to a temporary location specified by the ``tfrecord_dir``.</span>
<span class="sd">  The TensorFlow application will then be expected to read directly from this location during training.  However, if the input DataFrame was</span>
<span class="sd">  produced by the ``dfutil.loadTFRecords()`` method, i.e. originated from TFRecords on disk, then the `tfrecord_dir` will be set to the</span>
<span class="sd">  original source location of the TFRecords with the additional export step.</span>

<span class="sd">  Args:</span>
<span class="sd">    :train_fn: TensorFlow &quot;main&quot; function for training.</span>
<span class="sd">    :tf_args: Arguments specific to the TensorFlow &quot;main&quot; function.</span>
<span class="sd">    :export_fn: TensorFlow function for exporting a saved_model.</span>
<span class="sd">  &quot;&quot;&quot;</span>

  <span class="n">train_fn</span> <span class="o">=</span> <span class="kc">None</span>
  <span class="n">export_fn</span> <span class="o">=</span> <span class="kc">None</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">train_fn</span><span class="p">,</span> <span class="n">tf_args</span><span class="p">,</span> <span class="n">export_fn</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">TFEstimator</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">train_fn</span> <span class="o">=</span> <span class="n">train_fn</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">export_fn</span> <span class="o">=</span> <span class="n">export_fn</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">args</span> <span class="o">=</span> <span class="n">Namespace</span><span class="p">(</span><span class="n">tf_args</span><span class="p">)</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">_setDefault</span><span class="p">(</span><span class="n">input_mapping</span><span class="o">=</span><span class="p">{},</span>
                     <span class="n">cluster_size</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                     <span class="n">num_ps</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
                     <span class="n">driver_ps_nodes</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                     <span class="n">input_mode</span><span class="o">=</span><span class="n">TFCluster</span><span class="o">.</span><span class="n">InputMode</span><span class="o">.</span><span class="n">SPARK</span><span class="p">,</span>
                     <span class="n">protocol</span><span class="o">=</span><span class="s1">&#39;grpc&#39;</span><span class="p">,</span>
                     <span class="n">tensorboard</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                     <span class="n">model_dir</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                     <span class="n">export_dir</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                     <span class="n">tfrecord_dir</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                     <span class="n">batch_size</span><span class="o">=</span><span class="mi">100</span><span class="p">,</span>
                     <span class="n">epochs</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                     <span class="n">readers</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                     <span class="n">steps</span><span class="o">=</span><span class="mi">1000</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">_fit</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Trains a TensorFlow model and returns a TFModel instance with the same args/params pointing to a checkpoint or saved_model on disk.</span>

<span class="sd">    Args:</span>
<span class="sd">      :dataset: A Spark DataFrame with columns that will be mapped to TensorFlow tensors.</span>

<span class="sd">    Returns:</span>
<span class="sd">      A TFModel representing the trained model, backed on disk by a TensorFlow checkpoint or saved_model.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">sc</span> <span class="o">=</span> <span class="n">SparkContext</span><span class="o">.</span><span class="n">getOrCreate</span><span class="p">()</span>

    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;===== 1. train args: </span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">args</span><span class="p">))</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;===== 2. train params: </span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_paramMap</span><span class="p">))</span>
    <span class="n">local_args</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">merge_args_params</span><span class="p">()</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;===== 3. train args + params: </span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">local_args</span><span class="p">))</span>

    <span class="k">if</span> <span class="n">local_args</span><span class="o">.</span><span class="n">input_mode</span> <span class="o">==</span> <span class="n">TFCluster</span><span class="o">.</span><span class="n">InputMode</span><span class="o">.</span><span class="n">TENSORFLOW</span><span class="p">:</span>
      <span class="k">if</span> <span class="n">dfutil</span><span class="o">.</span><span class="n">isLoadedDF</span><span class="p">(</span><span class="n">dataset</span><span class="p">):</span>
        <span class="c1"># if just a DataFrame loaded from tfrecords, just point to original source path</span>
        <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Loaded DataFrame of TFRecord.&quot;</span><span class="p">)</span>
        <span class="n">local_args</span><span class="o">.</span><span class="n">tfrecord_dir</span> <span class="o">=</span> <span class="n">dfutil</span><span class="o">.</span><span class="n">loadedDF</span><span class="p">[</span><span class="n">dataset</span><span class="p">]</span>
      <span class="k">else</span><span class="p">:</span>
        <span class="c1"># otherwise, save as tfrecords and point to save path</span>
        <span class="k">assert</span> <span class="n">local_args</span><span class="o">.</span><span class="n">tfrecord_dir</span><span class="p">,</span> <span class="s2">&quot;Please specify --tfrecord_dir to export DataFrame to TFRecord.&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">getInputMapping</span><span class="p">():</span>
          <span class="c1"># if input mapping provided, filter only required columns before exporting</span>
          <span class="n">dataset</span> <span class="o">=</span> <span class="n">dataset</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="nb">list</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">getInputMapping</span><span class="p">()))</span>
        <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Exporting DataFrame </span><span class="si">{}</span><span class="s2"> as TFRecord to: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">dataset</span><span class="o">.</span><span class="n">dtypes</span><span class="p">,</span> <span class="n">local_args</span><span class="o">.</span><span class="n">tfrecord_dir</span><span class="p">))</span>
        <span class="n">dfutil</span><span class="o">.</span><span class="n">saveAsTFRecords</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">local_args</span><span class="o">.</span><span class="n">tfrecord_dir</span><span class="p">)</span>
        <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Done saving&quot;</span><span class="p">)</span>

    <span class="n">tf_args</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">argv</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">argv</span> <span class="k">else</span> <span class="n">local_args</span>
    <span class="n">cluster</span> <span class="o">=</span> <span class="n">TFCluster</span><span class="o">.</span><span class="n">run</span><span class="p">(</span><span class="n">sc</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_fn</span><span class="p">,</span> <span class="n">tf_args</span><span class="p">,</span> <span class="n">local_args</span><span class="o">.</span><span class="n">cluster_size</span><span class="p">,</span> <span class="n">local_args</span><span class="o">.</span><span class="n">num_ps</span><span class="p">,</span>
                            <span class="n">local_args</span><span class="o">.</span><span class="n">tensorboard</span><span class="p">,</span> <span class="n">local_args</span><span class="o">.</span><span class="n">input_mode</span><span class="p">,</span> <span class="n">driver_ps_nodes</span><span class="o">=</span><span class="n">local_args</span><span class="o">.</span><span class="n">driver_ps_nodes</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">local_args</span><span class="o">.</span><span class="n">input_mode</span> <span class="o">==</span> <span class="n">TFCluster</span><span class="o">.</span><span class="n">InputMode</span><span class="o">.</span><span class="n">SPARK</span><span class="p">:</span>
      <span class="c1"># feed data, using a deterministic order for input columns (lexicographic by key)</span>
      <span class="n">input_cols</span> <span class="o">=</span> <span class="nb">sorted</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">getInputMapping</span><span class="p">())</span>
      <span class="n">cluster</span><span class="o">.</span><span class="n">train</span><span class="p">(</span><span class="n">dataset</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="n">input_cols</span><span class="p">)</span><span class="o">.</span><span class="n">rdd</span><span class="p">,</span> <span class="n">local_args</span><span class="o">.</span><span class="n">epochs</span><span class="p">)</span>
    <span class="n">cluster</span><span class="o">.</span><span class="n">shutdown</span><span class="p">(</span><span class="n">grace_secs</span><span class="o">=</span><span class="mi">30</span><span class="p">)</span>

    <span class="c1"># Run export function, if provided</span>
    <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">export_fn</span><span class="p">:</span>
      <span class="k">assert</span> <span class="n">local_args</span><span class="o">.</span><span class="n">export_dir</span><span class="p">,</span> <span class="s2">&quot;Export function requires --export_dir to be set&quot;</span>
      <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Exporting saved_model (via export_fn) to: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">local_args</span><span class="o">.</span><span class="n">export_dir</span><span class="p">))</span>

      <span class="k">def</span> <span class="nf">_export</span><span class="p">(</span><span class="n">iterator</span><span class="p">,</span> <span class="n">fn</span><span class="p">,</span> <span class="n">args</span><span class="p">):</span>
        <span class="n">single_node_env</span><span class="p">(</span><span class="n">args</span><span class="p">)</span>
        <span class="n">fn</span><span class="p">(</span><span class="n">args</span><span class="p">)</span>

      <span class="c1"># Run on a single exeucutor</span>
      <span class="n">sc</span><span class="o">.</span><span class="n">parallelize</span><span class="p">([</span><span class="mi">1</span><span class="p">],</span> <span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">foreachPartition</span><span class="p">(</span><span class="k">lambda</span> <span class="n">it</span><span class="p">:</span> <span class="n">_export</span><span class="p">(</span><span class="n">it</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">export_fn</span><span class="p">,</span> <span class="n">tf_args</span><span class="p">))</span>

    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_copyValues</span><span class="p">(</span><span class="n">TFModel</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">args</span><span class="p">))</span></div>


<div class="viewcode-block" id="TFModel"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.TFModel">[docs]</a><span class="k">class</span> <span class="nc">TFModel</span><span class="p">(</span><span class="n">Model</span><span class="p">,</span> <span class="n">TFParams</span><span class="p">,</span>
              <span class="n">HasInputMapping</span><span class="p">,</span> <span class="n">HasOutputMapping</span><span class="p">,</span>
              <span class="n">HasBatchSize</span><span class="p">,</span>
              <span class="n">HasModelDir</span><span class="p">,</span> <span class="n">HasExportDir</span><span class="p">,</span> <span class="n">HasSignatureDefKey</span><span class="p">,</span> <span class="n">HasTagSet</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;Spark ML Model backed by a TensorFlow model checkpoint/saved_model on disk.</span>

<span class="sd">  During ``transform()``, each executor will run an independent, single-node instance of TensorFlow in parallel, so the model must fit in memory.</span>
<span class="sd">  The model/session will be loaded/initialized just once for each Spark Python worker, and the session will be cached for</span>
<span class="sd">  subsequent tasks/partitions to avoid re-loading the model for each partition.</span>

<span class="sd">  Args:</span>
<span class="sd">    :tf_args: Dictionary of arguments specific to TensorFlow &quot;main&quot; function.</span>
<span class="sd">  &quot;&quot;&quot;</span>

  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">tf_args</span><span class="p">):</span>
    <span class="nb">super</span><span class="p">(</span><span class="n">TFModel</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">args</span> <span class="o">=</span> <span class="n">Namespace</span><span class="p">(</span><span class="n">tf_args</span><span class="p">)</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">_setDefault</span><span class="p">(</span><span class="n">input_mapping</span><span class="o">=</span><span class="p">{},</span>
                     <span class="n">output_mapping</span><span class="o">=</span><span class="p">{},</span>
                     <span class="n">batch_size</span><span class="o">=</span><span class="mi">100</span><span class="p">,</span>
                     <span class="n">model_dir</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                     <span class="n">export_dir</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                     <span class="n">signature_def_key</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                     <span class="n">tag_set</span><span class="o">=</span><span class="kc">None</span><span class="p">)</span>

  <span class="k">def</span> <span class="nf">_transform</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Transforms the input DataFrame by applying the _run_model() mapPartitions function.</span>

<span class="sd">    Args:</span>
<span class="sd">      :dataset: A Spark DataFrame for TensorFlow inferencing.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">spark</span> <span class="o">=</span> <span class="n">SparkSession</span><span class="o">.</span><span class="n">builder</span><span class="o">.</span><span class="n">getOrCreate</span><span class="p">()</span>

    <span class="c1"># set a deterministic order for input/output columns (lexicographic by key)</span>
    <span class="n">input_cols</span> <span class="o">=</span> <span class="p">[</span><span class="n">col</span> <span class="k">for</span> <span class="n">col</span><span class="p">,</span> <span class="n">tensor</span> <span class="ow">in</span> <span class="nb">sorted</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">getInputMapping</span><span class="p">()</span><span class="o">.</span><span class="n">items</span><span class="p">())]</span>      <span class="c1"># input col =&gt; input tensor</span>
    <span class="n">output_cols</span> <span class="o">=</span> <span class="p">[</span><span class="n">col</span> <span class="k">for</span> <span class="n">tensor</span><span class="p">,</span> <span class="n">col</span> <span class="ow">in</span> <span class="nb">sorted</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">getOutputMapping</span><span class="p">()</span><span class="o">.</span><span class="n">items</span><span class="p">())]</span>    <span class="c1"># output tensor =&gt; output col</span>

    <span class="c1"># run single-node inferencing on each executor</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;input_cols: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">input_cols</span><span class="p">))</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;output_cols: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">output_cols</span><span class="p">))</span>

    <span class="c1"># merge args + params</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;===== 1. inference args: </span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">args</span><span class="p">))</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;===== 2. inference params: </span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_paramMap</span><span class="p">))</span>
    <span class="n">local_args</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">merge_args_params</span><span class="p">()</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;===== 3. inference args + params: </span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">local_args</span><span class="p">))</span>

    <span class="n">tf_args</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">argv</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">argv</span> <span class="k">else</span> <span class="n">local_args</span>
    <span class="n">rdd_out</span> <span class="o">=</span> <span class="n">dataset</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="n">input_cols</span><span class="p">)</span><span class="o">.</span><span class="n">rdd</span><span class="o">.</span><span class="n">mapPartitions</span><span class="p">(</span><span class="k">lambda</span> <span class="n">it</span><span class="p">:</span> <span class="n">_run_model</span><span class="p">(</span><span class="n">it</span><span class="p">,</span> <span class="n">local_args</span><span class="p">,</span> <span class="n">tf_args</span><span class="p">))</span>

    <span class="c1"># convert to a DataFrame-friendly format</span>
    <span class="n">rows_out</span> <span class="o">=</span> <span class="n">rdd_out</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="k">lambda</span> <span class="n">x</span><span class="p">:</span> <span class="n">Row</span><span class="p">(</span><span class="o">*</span><span class="n">x</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">spark</span><span class="o">.</span><span class="n">createDataFrame</span><span class="p">(</span><span class="n">rows_out</span><span class="p">,</span> <span class="n">output_cols</span><span class="p">)</span></div>


<span class="c1"># global to each python worker process on the executors</span>
<span class="n">global_sess</span> <span class="o">=</span> <span class="kc">None</span>            <span class="c1"># tf.Session cache</span>
<span class="n">global_args</span> <span class="o">=</span> <span class="kc">None</span>            <span class="c1"># args provided to the _run_model() method.  Any change will invalidate the global_sess cache.</span>


<span class="k">def</span> <span class="nf">_run_model</span><span class="p">(</span><span class="n">iterator</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="n">tf_args</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;mapPartitions function to run single-node inferencing from a checkpoint/saved_model, using the model&#39;s input/output mappings.</span>

<span class="sd">  Args:</span>
<span class="sd">    :iterator: input RDD partition iterator.</span>
<span class="sd">    :args: arguments for TFModel, in argparse format</span>
<span class="sd">    :tf_args: arguments for TensorFlow inferencing code, in argparse or ARGV format.</span>

<span class="sd">  Returns:</span>
<span class="sd">    An iterator of result data.</span>
<span class="sd">  &quot;&quot;&quot;</span>
  <span class="n">single_node_env</span><span class="p">(</span><span class="n">tf_args</span><span class="p">)</span>

  <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;===== input_mapping: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">args</span><span class="o">.</span><span class="n">input_mapping</span><span class="p">))</span>
  <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;===== output_mapping: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">args</span><span class="o">.</span><span class="n">output_mapping</span><span class="p">))</span>
  <span class="n">input_tensor_names</span> <span class="o">=</span> <span class="p">[</span><span class="n">tensor</span> <span class="k">for</span> <span class="n">col</span><span class="p">,</span> <span class="n">tensor</span> <span class="ow">in</span> <span class="nb">sorted</span><span class="p">(</span><span class="n">args</span><span class="o">.</span><span class="n">input_mapping</span><span class="o">.</span><span class="n">items</span><span class="p">())]</span>
  <span class="n">output_tensor_names</span> <span class="o">=</span> <span class="p">[</span><span class="n">tensor</span> <span class="k">for</span> <span class="n">tensor</span><span class="p">,</span> <span class="n">col</span> <span class="ow">in</span> <span class="nb">sorted</span><span class="p">(</span><span class="n">args</span><span class="o">.</span><span class="n">output_mapping</span><span class="o">.</span><span class="n">items</span><span class="p">())]</span>

  <span class="c1"># if using a signature_def_key, get input/output tensor info from the requested signature</span>
  <span class="k">if</span> <span class="n">args</span><span class="o">.</span><span class="n">signature_def_key</span><span class="p">:</span>
    <span class="k">assert</span> <span class="n">args</span><span class="o">.</span><span class="n">export_dir</span><span class="p">,</span> <span class="s2">&quot;Inferencing with signature_def_key requires --export_dir argument&quot;</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;===== loading meta_graph_def for tag_set (</span><span class="si">{0}</span><span class="s2">) from saved_model: </span><span class="si">{1}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">args</span><span class="o">.</span><span class="n">tag_set</span><span class="p">,</span> <span class="n">args</span><span class="o">.</span><span class="n">export_dir</span><span class="p">))</span>
    <span class="n">meta_graph_def</span> <span class="o">=</span> <span class="n">get_meta_graph_def</span><span class="p">(</span><span class="n">args</span><span class="o">.</span><span class="n">export_dir</span><span class="p">,</span> <span class="n">args</span><span class="o">.</span><span class="n">tag_set</span><span class="p">)</span>
    <span class="n">signature</span> <span class="o">=</span> <span class="n">signature_def_utils</span><span class="o">.</span><span class="n">get_signature_def_by_key</span><span class="p">(</span><span class="n">meta_graph_def</span><span class="p">,</span> <span class="n">args</span><span class="o">.</span><span class="n">signature_def_key</span><span class="p">)</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">debug</span><span class="p">(</span><span class="s2">&quot;signature: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">signature</span><span class="p">))</span>
    <span class="n">inputs_tensor_info</span> <span class="o">=</span> <span class="n">signature</span><span class="o">.</span><span class="n">inputs</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">debug</span><span class="p">(</span><span class="s2">&quot;inputs_tensor_info: </span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">inputs_tensor_info</span><span class="p">))</span>
    <span class="n">outputs_tensor_info</span> <span class="o">=</span> <span class="n">signature</span><span class="o">.</span><span class="n">outputs</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">debug</span><span class="p">(</span><span class="s2">&quot;outputs_tensor_info: </span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">outputs_tensor_info</span><span class="p">))</span>

  <span class="n">result</span> <span class="o">=</span> <span class="p">[]</span>

  <span class="k">global</span> <span class="n">global_sess</span><span class="p">,</span> <span class="n">global_args</span>
  <span class="k">if</span> <span class="n">global_sess</span> <span class="ow">and</span> <span class="n">global_args</span> <span class="o">==</span> <span class="n">args</span><span class="p">:</span>
    <span class="c1"># if graph/session already loaded/started (and using same args), just reuse it</span>
    <span class="n">sess</span> <span class="o">=</span> <span class="n">global_sess</span>
  <span class="k">else</span><span class="p">:</span>
    <span class="c1"># otherwise, create new session and load graph from disk</span>
    <span class="n">tf</span><span class="o">.</span><span class="n">reset_default_graph</span><span class="p">()</span>
    <span class="n">sess</span> <span class="o">=</span> <span class="n">tf</span><span class="o">.</span><span class="n">Session</span><span class="p">(</span><span class="n">graph</span><span class="o">=</span><span class="n">tf</span><span class="o">.</span><span class="n">get_default_graph</span><span class="p">())</span>
    <span class="k">if</span> <span class="n">args</span><span class="o">.</span><span class="n">export_dir</span><span class="p">:</span>
      <span class="k">assert</span> <span class="n">args</span><span class="o">.</span><span class="n">tag_set</span><span class="p">,</span> <span class="s2">&quot;Inferencing from a saved_model requires --tag_set&quot;</span>
      <span class="c1"># load graph from a saved_model</span>
      <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;===== restoring from saved_model: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">args</span><span class="o">.</span><span class="n">export_dir</span><span class="p">))</span>
      <span class="n">loader</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="n">sess</span><span class="p">,</span> <span class="n">args</span><span class="o">.</span><span class="n">tag_set</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s1">&#39;,&#39;</span><span class="p">),</span> <span class="n">args</span><span class="o">.</span><span class="n">export_dir</span><span class="p">)</span>
    <span class="k">elif</span> <span class="n">args</span><span class="o">.</span><span class="n">model_dir</span><span class="p">:</span>
      <span class="c1"># load graph from a checkpoint</span>
      <span class="n">ckpt</span> <span class="o">=</span> <span class="n">tf</span><span class="o">.</span><span class="n">train</span><span class="o">.</span><span class="n">latest_checkpoint</span><span class="p">(</span><span class="n">args</span><span class="o">.</span><span class="n">model_dir</span><span class="p">)</span>
      <span class="k">assert</span> <span class="n">ckpt</span><span class="p">,</span> <span class="s2">&quot;Invalid model checkpoint path: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">args</span><span class="o">.</span><span class="n">model_dir</span><span class="p">)</span>
      <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;===== restoring from checkpoint: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">ckpt</span> <span class="o">+</span> <span class="s2">&quot;.meta&quot;</span><span class="p">))</span>
      <span class="n">saver</span> <span class="o">=</span> <span class="n">tf</span><span class="o">.</span><span class="n">train</span><span class="o">.</span><span class="n">import_meta_graph</span><span class="p">(</span><span class="n">ckpt</span> <span class="o">+</span> <span class="s2">&quot;.meta&quot;</span><span class="p">,</span> <span class="n">clear_devices</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
      <span class="n">saver</span><span class="o">.</span><span class="n">restore</span><span class="p">(</span><span class="n">sess</span><span class="p">,</span> <span class="n">ckpt</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
      <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;Inferencing requires either --model_dir or --export_dir argument&quot;</span><span class="p">)</span>
    <span class="n">global_sess</span> <span class="o">=</span> <span class="n">sess</span>
    <span class="n">global_args</span> <span class="o">=</span> <span class="n">args</span>

  <span class="c1"># get list of input/output tensors (by name)</span>
  <span class="k">if</span> <span class="n">args</span><span class="o">.</span><span class="n">signature_def_key</span><span class="p">:</span>
    <span class="n">input_tensors</span> <span class="o">=</span> <span class="p">[</span><span class="n">inputs_tensor_info</span><span class="p">[</span><span class="n">t</span><span class="p">]</span><span class="o">.</span><span class="n">name</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="n">input_tensor_names</span><span class="p">]</span>
    <span class="n">output_tensors</span> <span class="o">=</span> <span class="p">[</span><span class="n">outputs_tensor_info</span><span class="p">[</span><span class="n">output_tensor_names</span><span class="p">[</span><span class="mi">0</span><span class="p">]]</span><span class="o">.</span><span class="n">name</span><span class="p">]</span>
  <span class="k">else</span><span class="p">:</span>
    <span class="n">input_tensors</span> <span class="o">=</span> <span class="p">[</span><span class="n">t</span> <span class="o">+</span> <span class="s1">&#39;:0&#39;</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="n">input_tensor_names</span><span class="p">]</span>
    <span class="n">output_tensors</span> <span class="o">=</span> <span class="p">[</span><span class="n">t</span> <span class="o">+</span> <span class="s1">&#39;:0&#39;</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="n">output_tensor_names</span><span class="p">]</span>

  <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;input_tensors: </span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">input_tensors</span><span class="p">))</span>
  <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;output_tensors: </span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">output_tensors</span><span class="p">))</span>

  <span class="c1"># feed data in batches and return output tensors</span>
  <span class="k">for</span> <span class="n">tensors</span> <span class="ow">in</span> <span class="n">yield_batch</span><span class="p">(</span><span class="n">iterator</span><span class="p">,</span> <span class="n">args</span><span class="o">.</span><span class="n">batch_size</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="n">input_tensor_names</span><span class="p">)):</span>
    <span class="n">inputs_feed_dict</span> <span class="o">=</span> <span class="p">{}</span>
    <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">input_tensors</span><span class="p">)):</span>
      <span class="n">inputs_feed_dict</span><span class="p">[</span><span class="n">input_tensors</span><span class="p">[</span><span class="n">i</span><span class="p">]]</span> <span class="o">=</span> <span class="n">tensors</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>

    <span class="n">outputs</span> <span class="o">=</span> <span class="n">sess</span><span class="o">.</span><span class="n">run</span><span class="p">(</span><span class="n">output_tensors</span><span class="p">,</span> <span class="n">feed_dict</span><span class="o">=</span><span class="n">inputs_feed_dict</span><span class="p">)</span>
    <span class="n">lengths</span> <span class="o">=</span> <span class="p">[</span><span class="nb">len</span><span class="p">(</span><span class="n">output</span><span class="p">)</span> <span class="k">for</span> <span class="n">output</span> <span class="ow">in</span> <span class="n">outputs</span><span class="p">]</span>
    <span class="n">input_size</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">tensors</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span>
    <span class="k">assert</span> <span class="nb">all</span><span class="p">([</span><span class="n">length</span> <span class="o">==</span> <span class="n">input_size</span> <span class="k">for</span> <span class="n">length</span> <span class="ow">in</span> <span class="n">lengths</span><span class="p">]),</span> <span class="s2">&quot;Output array sizes </span><span class="si">{}</span><span class="s2"> must match input size: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">lengths</span><span class="p">,</span> <span class="n">input_size</span><span class="p">)</span>
    <span class="n">python_outputs</span> <span class="o">=</span> <span class="p">[</span><span class="n">output</span><span class="o">.</span><span class="n">tolist</span><span class="p">()</span> <span class="k">for</span> <span class="n">output</span> <span class="ow">in</span> <span class="n">outputs</span><span class="p">]</span>      <span class="c1"># convert from numpy to standard python types</span>
    <span class="n">result</span><span class="o">.</span><span class="n">extend</span><span class="p">(</span><span class="nb">zip</span><span class="p">(</span><span class="o">*</span><span class="n">python_outputs</span><span class="p">))</span>                           <span class="c1"># convert to an array of tuples of &quot;output columns&quot;</span>

  <span class="k">return</span> <span class="n">result</span>


<div class="viewcode-block" id="single_node_env"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.single_node_env">[docs]</a><span class="k">def</span> <span class="nf">single_node_env</span><span class="p">(</span><span class="n">args</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;Sets up environment for a single-node TF session.</span>

<span class="sd">  Args:</span>
<span class="sd">    :args: command line arguments as either argparse args or argv list</span>
<span class="sd">  &quot;&quot;&quot;</span>
  <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
      <span class="n">sys</span><span class="o">.</span><span class="n">argv</span> <span class="o">=</span> <span class="n">args</span>
  <span class="k">elif</span> <span class="n">args</span><span class="o">.</span><span class="n">argv</span><span class="p">:</span>
      <span class="n">sys</span><span class="o">.</span><span class="n">argv</span> <span class="o">=</span> <span class="n">args</span><span class="o">.</span><span class="n">argv</span>

  <span class="c1"># ensure expanded CLASSPATH w/o glob characters (required for Spark 2.1 + JNI)</span>
  <span class="k">if</span> <span class="s1">&#39;HADOOP_PREFIX&#39;</span> <span class="ow">in</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span> <span class="ow">and</span> <span class="s1">&#39;TFOS_CLASSPATH_UPDATED&#39;</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">:</span>
      <span class="n">classpath</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s1">&#39;CLASSPATH&#39;</span><span class="p">]</span>
      <span class="n">hadoop_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s1">&#39;HADOOP_PREFIX&#39;</span><span class="p">],</span> <span class="s1">&#39;bin&#39;</span><span class="p">,</span> <span class="s1">&#39;hadoop&#39;</span><span class="p">)</span>
      <span class="n">hadoop_classpath</span> <span class="o">=</span> <span class="n">subprocess</span><span class="o">.</span><span class="n">check_output</span><span class="p">([</span><span class="n">hadoop_path</span><span class="p">,</span> <span class="s1">&#39;classpath&#39;</span><span class="p">,</span> <span class="s1">&#39;--glob&#39;</span><span class="p">])</span><span class="o">.</span><span class="n">decode</span><span class="p">()</span>
      <span class="n">logging</span><span class="o">.</span><span class="n">debug</span><span class="p">(</span><span class="s2">&quot;CLASSPATH: </span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">hadoop_classpath</span><span class="p">))</span>
      <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s1">&#39;CLASSPATH&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">classpath</span> <span class="o">+</span> <span class="n">os</span><span class="o">.</span><span class="n">pathsep</span> <span class="o">+</span> <span class="n">hadoop_classpath</span>
      <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s1">&#39;TFOS_CLASSPATH_UPDATED&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="s1">&#39;1&#39;</span>

  <span class="c1"># reserve GPU, if requested</span>
  <span class="k">if</span> <span class="n">tf</span><span class="o">.</span><span class="n">test</span><span class="o">.</span><span class="n">is_built_with_cuda</span><span class="p">():</span>
    <span class="c1"># GPU</span>
    <span class="n">num_gpus</span> <span class="o">=</span> <span class="n">args</span><span class="o">.</span><span class="n">num_gpus</span> <span class="k">if</span> <span class="s1">&#39;num_gpus&#39;</span> <span class="ow">in</span> <span class="n">args</span> <span class="k">else</span> <span class="mi">1</span>
    <span class="n">gpus_to_use</span> <span class="o">=</span> <span class="n">gpu_info</span><span class="o">.</span><span class="n">get_gpus</span><span class="p">(</span><span class="n">num_gpus</span><span class="p">)</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Using gpu(s): </span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">gpus_to_use</span><span class="p">))</span>
    <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s1">&#39;CUDA_VISIBLE_DEVICES&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">gpus_to_use</span>
    <span class="c1"># Note: if there is a GPU conflict (CUDA_ERROR_INVALID_DEVICE), the entire task will fail and retry.</span>
  <span class="k">else</span><span class="p">:</span>
    <span class="c1"># CPU</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Using CPU&quot;</span><span class="p">)</span>
    <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s1">&#39;CUDA_VISIBLE_DEVICES&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="s1">&#39;&#39;</span></div>


<div class="viewcode-block" id="get_meta_graph_def"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.get_meta_graph_def">[docs]</a><span class="k">def</span> <span class="nf">get_meta_graph_def</span><span class="p">(</span><span class="n">saved_model_dir</span><span class="p">,</span> <span class="n">tag_set</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;Utility function to read a meta_graph_def from disk.</span>

<span class="sd">  From `saved_model_cli.py &lt;https://github.com/tensorflow/tensorflow/blob/8e0e8d41a3a8f2d4a6100c2ea1dc9d6c6c4ad382/tensorflow/python/tools/saved_model_cli.py#L186&gt;`_</span>

<span class="sd">  Args:</span>
<span class="sd">    :saved_model_dir: path to saved_model.</span>
<span class="sd">    :tag_set: list of string tags identifying the TensorFlow graph within the saved_model.</span>

<span class="sd">  Returns:</span>
<span class="sd">    A TensorFlow meta_graph_def, or raises an Exception otherwise.</span>
<span class="sd">  &quot;&quot;&quot;</span>
  <span class="n">saved_model</span> <span class="o">=</span> <span class="n">reader</span><span class="o">.</span><span class="n">read_saved_model</span><span class="p">(</span><span class="n">saved_model_dir</span><span class="p">)</span>
  <span class="n">set_of_tags</span> <span class="o">=</span> <span class="nb">set</span><span class="p">(</span><span class="n">tag_set</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s1">&#39;,&#39;</span><span class="p">))</span>
  <span class="k">for</span> <span class="n">meta_graph_def</span> <span class="ow">in</span> <span class="n">saved_model</span><span class="o">.</span><span class="n">meta_graphs</span><span class="p">:</span>
    <span class="k">if</span> <span class="nb">set</span><span class="p">(</span><span class="n">meta_graph_def</span><span class="o">.</span><span class="n">meta_info_def</span><span class="o">.</span><span class="n">tags</span><span class="p">)</span> <span class="o">==</span> <span class="n">set_of_tags</span><span class="p">:</span>
      <span class="k">return</span> <span class="n">meta_graph_def</span>
  <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;MetaGraphDef associated with tag-set </span><span class="si">{0}</span><span class="s2"> could not be found in SavedModel&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">tag_set</span><span class="p">))</span></div>


<div class="viewcode-block" id="yield_batch"><a class="viewcode-back" href="../../tensorflowonspark.pipeline.html#tensorflowonspark.pipeline.yield_batch">[docs]</a><span class="k">def</span> <span class="nf">yield_batch</span><span class="p">(</span><span class="n">iterable</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">,</span> <span class="n">num_tensors</span><span class="o">=</span><span class="mi">1</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;Generator that yields batches of a DataFrame iterator.</span>

<span class="sd">  Args:</span>
<span class="sd">    :iterable: Spark partition iterator.</span>
<span class="sd">    :batch_size: number of items to retrieve per invocation.</span>
<span class="sd">    :num_tensors: number of tensors (columns) expected in each item.</span>

<span class="sd">  Returns:</span>
<span class="sd">    An array of ``num_tensors`` arrays, each of length `batch_size`</span>
<span class="sd">  &quot;&quot;&quot;</span>
  <span class="n">tensors</span> <span class="o">=</span> <span class="p">[[]</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_tensors</span><span class="p">)]</span>
  <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="n">iterable</span><span class="p">:</span>
    <span class="k">if</span> <span class="n">item</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
      <span class="k">break</span>
    <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_tensors</span><span class="p">):</span>
      <span class="n">tmp</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="n">item</span><span class="p">[</span><span class="n">i</span><span class="p">])</span> <span class="k">if</span> <span class="nb">type</span><span class="p">(</span><span class="n">item</span><span class="p">[</span><span class="n">i</span><span class="p">])</span> <span class="ow">is</span> <span class="nb">bytearray</span> <span class="k">else</span> <span class="n">item</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>
      <span class="n">tensors</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">tmp</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">tensors</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="o">&gt;=</span> <span class="n">batch_size</span><span class="p">:</span>
      <span class="k">yield</span> <span class="n">tensors</span>
      <span class="n">tensors</span> <span class="o">=</span> <span class="p">[[]</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_tensors</span><span class="p">)]</span>
  <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">tensors</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
      <span class="k">yield</span> <span class="n">tensors</span></div>
</pre></div>

          </div>
        </div>
      </div>
      <div class="sphinxsidebar" role="navigation" aria-label="main navigation">
        <div class="sphinxsidebarwrapper">
<div id="searchbox" style="display: none" role="search">
  <h3>Quick search</h3>
    <div class="searchformwrapper">
    <form class="search" action="../../search.html" method="get">
      <input type="text" name="q" />
      <input type="submit" value="Go" />
      <input type="hidden" name="check_keywords" value="yes" />
      <input type="hidden" name="area" value="default" />
    </form>
    </div>
</div>
<script type="text/javascript">$('#searchbox').show(0);</script>
        </div>
      </div>
      <div class="clearer"></div>
    </div>
    <div class="related" role="navigation" aria-label="related navigation">
      <h3>Navigation</h3>
      <ul>
        <li class="right" style="margin-right: 10px">
          <a href="../../genindex.html" title="General Index"
             >index</a></li>
        <li class="right" >
          <a href="../../py-modindex.html" title="Python Module Index"
             >modules</a> |</li>
        <li class="nav-item nav-item-0"><a href="../../index.html">TensorFlowOnSpark 1.3.3 documentation</a> &#187;</li>
          <li class="nav-item nav-item-1"><a href="../index.html" >Module code</a> &#187;</li> 
      </ul>
    </div>
    <div class="footer" role="contentinfo">
        &#169; Copyright 2018, Yahoo Inc.
      Created using <a href="http://sphinx-doc.org/">Sphinx</a> 1.7.9.
    </div>
  </body>
</html>